import one_hot
import init
import torch
from dataset import _dataset
from torch.utils.data import DataLoader

if __name__ == '__main__':
    test_dataset = _dataset("./datasets/test/")
    test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=True)
    _model = torch.load("model.pth").cuda()

    correct = 0
    test_len = test_dataset.__len__()
    # print(test_len)
    for i, (img, lab) in enumerate(test_dataloader):
        img = img.cuda()
        lab = lab.cuda()
        lab = lab.view(-1, len(init.captcha_array))
        # print(lab.shape)
        lab_text = one_hot.vec2Text(lab)
        output = _model(img).cuda()
        output = output.view(-1, len(init.captcha_array))
        output_text = one_hot.vec2Text(output)
        # print(output_text)
        if(lab_text == output_text):
            correct += 1
            print("正确值:{}, 预测值{}".format(lab_text, output_text))

    print("正确率:{}%".format(correct / test_len * 100))

        